import os
import matplotlib.pyplot as plt
import csv
from cectest.common.alg import *
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

"""创建结果文件夹"""
def setup_results_directory(results_dir):
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)
"""保存单个函数的收敛曲线"""
def save_convergence_plot(results_dir, func_num, conv_curve):
    plt.figure()
    plt.plot(conv_curve, label=f"F{func_num} Convergence")
    plt.xlabel("Iterations")
    plt.ylabel("Best Fitness")
    plt.title(f"Convergence Curve for Function F{func_num}")
    plt.legend()
    # plot_path_png = os.path.join(results_dir, f"F{func_num}_convergence.png")
    # plt.savefig(plot_path_png)
    plot_path_svg = os.path.join(results_dir, f"F{func_num}_convergence.svg")
    plt.savefig(plot_path_svg)
    plt.close()
"""绘制并保存所有函数的收敛曲线"""
def save_all_convergence_curves(results_dir, conv_curves):
    plt.figure(figsize=(15, 8))
    for func_num in range(1, 31):
        plt.plot(conv_curves[func_num], label=f'F{func_num}', alpha=0.7)
    plt.xlabel("Iterations")
    plt.ylabel("Best Fitness (log scale)")
    plt.title("All Convergence Curves")
    plt.yscale('log')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', ncol=2)
    plt.tight_layout()
    # plot_all_path_png = os.path.join(results_dir, "all_convergence_curves.png")
    # plt.savefig(plot_all_path_png, bbox_inches='tight')
    plot_all_path_svg = os.path.join(results_dir, "all_convergence_curves.svg")
    plt.savefig(plot_all_path_svg, bbox_inches='tight')
    plt.close()
"""将每个函数对应的最优适应度值保存到 CSV 文件中"""
def save_fitness_results(results_dir, fun_fitness, mean_fitness, std_fitness):
    csv_path = os.path.join(results_dir, "final_fitness_values.csv")
    with open(csv_path, mode="w", newline="") as file:
        writer = csv.writer(file)
        writer.writerow(["Function", "Best Fitness", "mean_fitness", "std_fitness"])
        for func_num in range(1, 31):
            writer.writerow([f"F{func_num}", fun_fitness[func_num], mean_fitness[func_num], std_fitness[func_num]])

'''将不同算法在同一个函数（F1-F30）上的迭代曲线绘制到一张图中，每张图对应一个函数，并保存为 PNG 和 SVG 格式'''
def plot_algorithms_by_function(algorithms, results_dir):
    # 创建结果目录
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)

    # 遍历每个函数（F1-F30）
    for func_num in range(1, 31):
        plt.figure(figsize=(10, 6))

        # 绘制每种算法在该函数上的收敛曲线
        for algo_name, conv_curves in algorithms.items():
            plt.plot(conv_curves[func_num], label=f'{algo_name}', alpha=0.7)

        # 设置图形属性
        plt.xlabel("Iterations")
        plt.ylabel("Best Fitness (log scale)")
        plt.title(f"Convergence Curves for F{func_num}")
        plt.yscale('log')  # 使用对数坐标
        plt.legend()  # 显示图例
        plt.tight_layout()

        # 保存为 PNG 格式
        # plot_path_png = os.path.join(results_dir, f"F{func_num}_convergence.png")
        # plt.savefig(plot_path_png, bbox_inches='tight')

        # 保存为 SVG 格式
        plot_path_svg = os.path.join(results_dir, f"F{func_num}_convergence.svg")
        plt.savefig(plot_path_svg, bbox_inches='tight')

        plt.close()

    print(f"All function convergence curves have been saved in the '{results_dir}' folder.")
'''3d曲面图'''
def visualize_function(benchmark_func, func_num, lb, ub, dim, results_dir):
    x = np.linspace(lb, ub, 100)
    y = np.linspace(lb, ub, 100)
    X, Y = np.meshgrid(x, y)

    # For each (x, y) pair, create the corresponding input to the function
    Z = np.array([benchmark_func(np.array([[xi, yi]])) for xi, yi in zip(np.ravel(X), np.ravel(Y))])

    # Reshape Z to match X and Y for plotting
    Z = Z.reshape(X.shape)

    # Plotting the 3D surface
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_surface(X, Y, Z, cmap='viridis')
    ax.set_title(f"Function F{func_num} Visualization")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("f(x, y)")
    plot_path_svg = os.path.join(results_dir, f"F{func_num}_3d.svg")
    plt.savefig(plot_path_svg)
    plt.close()

# dim纬度 F1-F30均可取值2, 10, 20, 30, 50 or 100
def run_algorithm_experiments(algorithm_func, results_dir, dim, pop_size, Tmax, ub, lb, population):
    """运行实验并保存结果"""
    setup_results_directory(results_dir)
    fun_fitness = {} # 记录cec测试集上每个函数的最优值
    conv_curves = {} # 记录cec测试集上每个函数的迭代值
    mean_fitness = {} # 记录cec测试集上每个函数的平均值
    std_fitness = {} # 记录cec测试集上每个函数的标准差

    for func_num in range(1, 31):
        print(f"Running Function F{func_num}...")
        best_fit, best_pos, conv_curve = algorithm_func(population, pop_size, Tmax, ub, lb, dim, func_num)
        # if func_num < 11 : visualize_function(all_functions[func_num - 1], func_num, lb, ub, dim, results_dir)
        save_convergence_plot(results_dir, func_num, conv_curve)
        fun_fitness[func_num] = best_fit
        conv_curves[func_num] = conv_curve
        mean_fitness[func_num] = np.mean(conv_curve)
        std_fitness[func_num] = np.std(conv_curve)
        print(f"Function F{func_num} Final Best Fitness: {np.min(fun_fitness[func_num]):.10e}")
        print(f"Function F{func_num}: Best Fitness = {np.min(fun_fitness[func_num]):.2e}, Mean Fitness = {mean_fitness[func_num]:.2e}, Std Dev = {std_fitness[func_num]:.2e}")

    save_all_convergence_curves(results_dir, conv_curves)
    save_fitness_results(results_dir, fun_fitness, mean_fitness, std_fitness)
    print(f"All results have been saved in the '{results_dir}' folder.")
    return conv_curves